"""
RAG 系统核心逻辑
整合查询分类、策略选择、向量检索和答案生成
"""

import json
from typing import List, Dict, Any, Optional
from base.config import config
from base.logger import logger
from .query_classifier import QueryClassifier
from .strategy_selector import StrategySelector
from .vector_store import VectorStore
from .prompts import RAG_PROMPT_TEMPLATE, MULTI_TURN_TEMPLATE


class RAGSystem:
    """RAG 系统核心"""
    
    def __init__(self):
        self.query_classifier = QueryClassifier()
        self.strategy_selector = StrategySelector()
        self.vector_store = VectorStore()
        self.chat_history = []
        self.max_history_length = config.getint('rag', 'max_history_length', fallback=10)
    
    def add_documents(self, documents: List[Dict[str, Any]]) -> bool:
        """
        添加文档到知识库
        
        Args:
            documents: 文档列表，每个文档包含 doc_id, content, embedding, metadata
            
        Returns:
            是否添加成功
        """
        try:
            logger.info(f"开始添加 {len(documents)} 个文档到知识库")
            
            # 验证文档格式
            for doc in documents:
                if not all(key in doc for key in ['doc_id', 'content', 'embedding']):
                    raise ValueError(f"文档格式不正确: {doc}")
            
            # 插入到向量数据库
            success = self.vector_store.insert_documents(documents)
            
            if success:
                logger.info("文档添加成功")
            else:
                logger.error("文档添加失败")
            
            return success
            
        except Exception as e:
            logger.error(f"添加文档时发生错误: {e}")
            return False
    
    def query(self, question: str, use_history: bool = True) -> Dict[str, Any]:
        """
        执行 RAG 查询
        
        Args:
            question: 用户问题
            use_history: 是否使用对话历史
            
        Returns:
            查询结果
        """
        try:
            logger.info(f"开始处理查询: {question}")
            
            # 1. 查询分类和特征提取
            query_features = self.query_classifier.get_query_features(question)
            logger.info(f"查询特征: {query_features}")
            
            # 2. 选择检索策略
            strategy = self.strategy_selector.select_strategy(question)
            logger.info(f"检索策略: {strategy}")
            
            # 3. 生成查询向量（这里需要实际的嵌入模型）
            query_embedding = self._generate_query_embedding(question)
            
            # 4. 向量检索
            similar_docs = self.vector_store.search_similar(
                query_embedding=query_embedding,
                top_k=strategy['top_k'],
                similarity_threshold=strategy['similarity_threshold']
            )
            
            # 5. 构建上下文
            context = self._build_context(similar_docs)
            
            # 6. 生成答案
            answer = self._generate_answer(question, context, use_history)
            
            # 7. 更新对话历史
            if use_history:
                self._update_chat_history(question, answer)
            
            # 8. 构建返回结果
            result = {
                'question': question,
                'answer': answer,
                'context': context,
                'similar_docs': similar_docs,
                'query_features': query_features,
                'strategy': strategy,
                'success': True
            }
            
            logger.info("查询处理完成")
            return result
            
        except Exception as e:
            logger.error(f"查询处理失败: {e}")
            return {
                'question': question,
                'answer': f"抱歉，处理您的问题时发生了错误: {str(e)}",
                'success': False
            }
    
    def _generate_query_embedding(self, query: str) -> List[float]:
        """
        生成查询向量
        注意：这里需要实际的嵌入模型实现
        
        Args:
            query: 查询文本
            
        Returns:
            查询向量
        """
        # 这里应该使用实际的嵌入模型，比如 sentence-transformers
        # 目前返回随机向量作为示例
        import random
        dimension = self.vector_store.dimension
        embedding = [random.random() for _ in range(dimension)]
        
        logger.warning("使用随机向量作为查询嵌入，请替换为实际的嵌入模型")
        return embedding
    
    def _build_context(self, similar_docs: List[Dict[str, Any]]) -> str:
        """
        构建上下文
        
        Args:
            similar_docs: 相似文档列表
            
        Returns:
            上下文字符串
        """
        if not similar_docs:
            return "没有找到相关的文档内容。"
        
        context_parts = []
        for i, doc in enumerate(similar_docs, 1):
            context_part = f"文档 {i}:\n{doc['content']}\n"
            context_parts.append(context_part)
        
        return "\n".join(context_parts)
    
    def _generate_answer(self, question: str, context: str, use_history: bool = True) -> str:
        """
        生成答案
        注意：这里需要实际的 LLM 实现
        
        Args:
            question: 用户问题
            context: 上下文
            use_history: 是否使用历史
            
        Returns:
            生成的答案
        """
        # 选择提示模板
        if use_history and self.chat_history:
            prompt_template = MULTI_TURN_TEMPLATE
            chat_history_str = self._format_chat_history()
            prompt = prompt_template.format(
                chat_history=chat_history_str,
                context=context,
                question=question
            )
        else:
            prompt_template = RAG_PROMPT_TEMPLATE
            prompt = prompt_template.format(
                context=context,
                question=question
            )
        
        # 这里应该调用实际的 LLM API
        # 目前返回基于上下文的简单回答
        if "没有找到相关的文档内容" in context:
            answer = "抱歉，我没有找到相关的信息来回答您的问题。请尝试重新表述您的问题或提供更多详细信息。"
        else:
            answer = f"根据相关文档，{question}的回答如下：\n\n基于提供的文档内容，我可以为您提供相关信息。但请注意，这是一个简化的回答，实际应用中需要集成真正的大语言模型来生成更准确和详细的答案。"
        
        logger.warning("使用简化的答案生成，请替换为实际的 LLM 调用")
        return answer
    
    def _update_chat_history(self, question: str, answer: str):
        """更新对话历史"""
        self.chat_history.append({
            'question': question,
            'answer': answer,
            'timestamp': self._get_current_timestamp()
        })
        
        # 保持历史长度限制
        if len(self.chat_history) > self.max_history_length:
            self.chat_history = self.chat_history[-self.max_history_length:]
    
    def _format_chat_history(self) -> str:
        """格式化对话历史"""
        if not self.chat_history:
            return "无对话历史"
        
        history_parts = []
        for i, turn in enumerate(self.chat_history[-5:], 1):  # 只使用最近5轮对话
            history_part = f"第{i}轮:\n用户: {turn['question']}\n助手: {turn['answer']}\n"
            history_parts.append(history_part)
        
        return "\n".join(history_parts)
    
    def _get_current_timestamp(self) -> str:
        """获取当前时间戳"""
        from datetime import datetime
        return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    def clear_history(self):
        """清空对话历史"""
        self.chat_history = []
        logger.info("对话历史已清空")
    
    def get_system_stats(self) -> Dict[str, Any]:
        """获取系统统计信息"""
        try:
            vector_stats = self.vector_store.get_collection_stats()
            
            stats = {
                'vector_store': vector_stats,
                'chat_history_length': len(self.chat_history),
                'max_history_length': self.max_history_length,
                'system_status': 'running'
            }
            
            return stats
            
        except Exception as e:
            logger.error(f"获取系统统计信息失败: {e}")
            return {'system_status': 'error', 'error': str(e)}